package com.liyuxiang.graph.node;

import com.alibaba.cloud.ai.graph.NodeOutput;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.streaming.StreamingChatGenerator;
import com.liyuxiang.graph.service.DiagnosisContextService;
import com.liyuxiang.graph.context.DiagnosisContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.PromptTemplate;
import reactor.core.publisher.Flux;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;

public class StreamSymptomDetailNode implements NodeAction {
    
    private static final Logger logger = LoggerFactory.getLogger(StreamSymptomDetailNode.class);
    
    private final ChatClient chatClient;
    private final DiagnosisContextService contextService;
    
    private static final String SYMPTOM_DETAIL_TEMPLATE = """
            你是专业的智能问诊助手，正在收集用户的症状详情以便准确推荐科室。
            
            用户输入: {query}
            对话历史: {conversationHistory}
            当前问题轮次: {questionCount}
            已收集症状: {currentSymptoms}
            
            请分析用户提供的症状信息，并温暖专业地询问必要的细节。
            
            询问策略：
            1. 针对用户描述的症状，询问具体的细节（时间、程度、位置、诱因等）
            2. 根据症状特点，询问相关的伴随症状
            3. 询问既往病史和用药情况（如果相关）
            4. 每次只问1-2个重点问题，避免过多询问
            5. 保持温暖关怀的语调，体现专业性
            
            分析要点：
            - 评估已收集信息的完整性（症状描述、时间、程度、位置、诱因、伴随症状）
            - 判断是否已收集足够信息进行科室推荐
            - 识别关键的缺失信息
            - 当信息足够完整时（completenessScore >= 0.8）或问题轮次>=3时，建议结束询问
            
            请返回JSON格式：
            {{
              "response": "温暖专业的询问内容或推荐科室的建议",
              "symptomsCollected": {{
                "symptom1": "详细描述",
                "symptom2": "详细描述"
              }},
              "questionCount": {questionCount},
              "shouldContinue": false,
              "completenessScore": 0.8,
              "nextFocus": "下一步重点询问的方向或科室推荐理由",
              "medicalRelevance": "high",
              "readyForRecommendation": true,
              "reasoning": "基于收集的信息进行分析的理由"
            }}
            
            注意：
            - 如果completenessScore >= 0.8 或 questionCount >= 3，设置shouldContinue=false, readyForRecommendation=true
            - 如果信息已足够，在response中可以直接给出科室建议
            - reasoning字段说明为什么认为信息已够或需要继续收集
            """;
    
    public StreamSymptomDetailNode(ChatClient.Builder chatClientBuilder, DiagnosisContextService contextService) {
        this.chatClient = chatClientBuilder.build();
        this.contextService = contextService;
    }
    
    @Override
    public Map<String, Object> apply(OverAllState state) throws Exception {
        String query = state.value("query", "");
        String userId = state.value("userId", "unknown");
        String threadId = state.value("threadId", "unknown");
        Integer questionCount = state.value("questionCount", 0);
        Map<String, Object> currentSymptoms = state.value("symptoms", new HashMap<>());
        
        // 获取实时对话历史
        final List<String> realTimeHistory;
        final Integer finalQuestionCount;
        final Map<String, Object> finalCurrentSymptoms;
        
        Optional<DiagnosisContext> contextOpt = contextService.getContext(threadId);
        if (contextOpt.isPresent()) {
            realTimeHistory = contextOpt.get().getConversationHistory();
            finalQuestionCount = contextOpt.get().getQuestionCount();
            finalCurrentSymptoms = contextOpt.get().getSymptoms();
        } else {
            realTimeHistory = List.of();
            finalQuestionCount = questionCount;
            finalCurrentSymptoms = currentSymptoms;
        }
        
        logger.info("开始流式症状详情收集 - 用户ID: {}, 问题轮次: {}, 已收集症状数: {}", 
                userId, questionCount, currentSymptoms.size());
        
        // 构建对话上下文
        String conversationContext = buildConversationContext(realTimeHistory);
        String symptomsContext = buildSymptomsContext(finalCurrentSymptoms);
        
        // 创建流式AI调用 - 正确使用PromptTemplate
        Map<String, Object> templateParams = new HashMap<>();
        templateParams.put("query", query);
        templateParams.put("conversationHistory", conversationContext);
        templateParams.put("questionCount", String.valueOf(finalQuestionCount));
        templateParams.put("currentSymptoms", symptomsContext);
        
        PromptTemplate promptTemplate = new PromptTemplate(SYMPTOM_DETAIL_TEMPLATE);
        String renderedPrompt = promptTemplate.render(templateParams);
        
        Flux<ChatResponse> chatResponseFlux = this.chatClient.prompt()
                .user(renderedPrompt)
                .stream()
                .chatResponse();
        
        // 创建流式生成器
        AsyncGenerator<? extends NodeOutput> generator = StreamingChatGenerator.builder()
                                                                               .startingNode("symptomDetailStream")
                                                                               .startingState(state)
                                                                               .mapResult(response -> {
                    String text = response.getResult().getOutput().getText();
                    logger.info("症状详情收集流式响应: {}", text);
                    
                    // 解析AI响应
                    final String finalQuery = query;
                    Map<String, Object> symptomResult = parseSymptomResponse(text, finalQuery, realTimeHistory, finalQuestionCount, finalCurrentSymptoms);
                    
                    return symptomResult;
                })
                                                                               .build(chatResponseFlux);
        
        // 返回流式生成器
        Map<String, Object> result = new HashMap<>();
        result.put("symptomDetailAnalysis", generator);
        result.put("questionCount", finalQuestionCount + 1);
        result.put("conversationHistory", realTimeHistory);
        
        logger.info("流式症状详情收集节点配置完成 - 用户ID: {}", userId);
        
        return result;
    }
    
    private String buildConversationContext(List<String> conversationHistory) {
        if (conversationHistory == null || conversationHistory.isEmpty()) {
            return "无对话历史";
        }
        
        // 取最后5轮对话作为上下文
        int start = Math.max(0, conversationHistory.size() - 10);
        List<String> recentHistory = conversationHistory.subList(start, conversationHistory.size());
        
        StringBuilder context = new StringBuilder();
        for (String message : recentHistory) {
            context.append(message).append("\n");
        }
        
        return context.toString().trim();
    }
    
    private String buildSymptomsContext(Map<String, Object> symptoms) {
        if (symptoms == null || symptoms.isEmpty()) {
            return "暂无收集到的症状信息";
        }
        
        StringBuilder context = new StringBuilder();
        symptoms.forEach((key, value) -> 
            context.append(key).append(": ").append(value).append("\n"));
        
        return context.toString().trim();
    }
    
    @SuppressWarnings("unchecked")
    private Map<String, Object> parseSymptomResponse(String response, String originalQuery, 
                                                   List<String> history, Integer questionCount, 
                                                   Map<String, Object> currentSymptoms) {
        Map<String, Object> result = new HashMap<>();
        
        try {
            // 尝试解析JSON响应
            if (response.contains("{") && response.contains("}")) {
                int startIndex = response.indexOf('{');
                int endIndex = response.lastIndexOf('}') + 1;
                String jsonStr = response.substring(startIndex, endIndex);
                
                // 提取AI回复
                String aiResponse = extractJsonValue(jsonStr, "response", "请继续描述您的症状细节");
                result.put("response", aiResponse);
                
                // 更新症状信息
                Map<String, Object> updatedSymptoms = new HashMap<>(currentSymptoms);
                String symptomsStr = extractJsonValue(jsonStr, "symptomsCollected", "{}");
                if (!symptomsStr.equals("{}")) {
                    // 简化处理，实际应该解析JSON
                    updatedSymptoms.put("latest_symptom", originalQuery);
                }
                result.put("symptoms", updatedSymptoms);
                
                // 更新问题计数
                result.put("questionCount", questionCount + 1);
                
                // 完整性评分
                String completenessStr = extractJsonValue(jsonStr, "completenessScore", "0.5");
                double completenessScore = Double.parseDouble(completenessStr);
                result.put("completenessScore", completenessScore);
                
                // 判断是否继续询问
                String shouldContinueStr = extractJsonValue(jsonStr, "shouldContinue", "true");
                String readyForRecommendationStr = extractJsonValue(jsonStr, "readyForRecommendation", "false");
                boolean readyForRecommendation = Boolean.parseBoolean(readyForRecommendationStr);
                
                // 智能判断：基于完整性评分、轮次和AI建议
                boolean shouldContinue = Boolean.parseBoolean(shouldContinueStr) 
                    && questionCount < 4  // 最多3轮问诊
                    && completenessScore < 0.8 
                    && !readyForRecommendation;
                    
                result.put("shouldContinue", shouldContinue);
                result.put("readyForRecommendation", readyForRecommendation);
                
                // 推理信息
                String reasoning = extractJsonValue(jsonStr, "reasoning", "继续收集症状详情");
                result.put("reasoning", reasoning);
                
                result.put("fullAiAnalysis", true);
                
            } else {
                // 文本分析备用方案
                String fallbackResponse = generateFallbackResponse(questionCount, currentSymptoms);
                result.put("response", fallbackResponse);
                result.put("symptoms", currentSymptoms);
                result.put("questionCount", questionCount + 1);
                
                // 智能的备用判断逻辑
                boolean shouldContinue = questionCount < 3 && !isSymptomInformationSufficient(currentSymptoms, questionCount);
                result.put("shouldContinue", shouldContinue);
                result.put("readyForRecommendation", !shouldContinue);
                result.put("completenessScore", calculateCompleteness(currentSymptoms, questionCount));
                result.put("reasoning", shouldContinue ? "需要更多症状细节" : "收集的信息已基本足够，可以推荐科室");
                result.put("fullAiAnalysis", false);
            }
            
            // 更新对话历史
            List<String> updatedHistory = new java.util.ArrayList<>(history);
            String currentUserInput = "用户: " + originalQuery;
            if (!updatedHistory.contains(currentUserInput)) {
                updatedHistory.add(currentUserInput);
            }
            
            // 添加AI回复到历史
            if (result.containsKey("response")) {
                String aiResponse = (String) result.get("response");
                updatedHistory.add("助手: " + aiResponse);
            }
            
            result.put("conversationHistory", updatedHistory);
            
        } catch (Exception e) {
            logger.error("解析症状详情响应失败: {}", e.getMessage());
            
            // 提供默认值
            result.put("response", "请继续详细描述您的症状，包括出现时间、严重程度等");
            result.put("symptoms", currentSymptoms);
            result.put("questionCount", questionCount + 1);
            result.put("shouldContinue", true);
            result.put("completenessScore", 0.5);
            result.put("fullAiAnalysis", false);
            result.put("error", e.getMessage());
        }
        
        return result;
    }
    
    private String extractJsonValue(String json, String key, String defaultValue) {
        try {
            String pattern = "\"" + key + "\"\\s*:\\s*\"([^\"]+)\"";
            java.util.regex.Pattern p = java.util.regex.Pattern.compile(pattern);
            java.util.regex.Matcher m = p.matcher(json);
            if (m.find()) {
                return m.group(1);
            }
        } catch (Exception e) {
            logger.warn("提取JSON值失败 - key: {}, 错误: {}", key, e.getMessage());
        }
        return defaultValue;
    }
    
    private String generateFallbackResponse(Integer questionCount, Map<String, Object> currentSymptoms) {
        if (questionCount == 1) {
            return "感谢您的描述。请问这些症状是什么时候开始出现的？严重程度如何？";
        } else if (questionCount == 2) {
            return "请问还有其他伴随的症状吗？比如发热、头痛等？";
        } else if (questionCount == 3) {
            return "请问您之前有类似的情况吗？最近有没有服用什么药物？";
        } else {
            return "根据您提供的信息，我来为您推荐合适的科室。";
        }
    }
    
    /**
     * 判断症状信息是否足够充分
     */
    private boolean isSymptomInformationSufficient(Map<String, Object> symptoms, Integer questionCount) {
        // 如果已经问了3轮或更多，认为信息足够
        if (questionCount >= 3) {
            return true;
        }
        
        // 如果症状信息很丰富，也认为足够
        if (symptoms != null && symptoms.size() >= 3) {
            return true;
        }
        
        // 其他情况认为信息不够
        return false;
    }
    
    /**
     * 计算症状收集的完整性
     */
    private double calculateCompleteness(Map<String, Object> symptoms, Integer questionCount) {
        double score = 0.0;
        
        // 基础分数：基于问诊轮次
        score += Math.min(questionCount * 0.25, 0.75);
        
        // 症状信息丰富度
        if (symptoms != null) {
            score += Math.min(symptoms.size() * 0.1, 0.25);
        }
        
        return Math.min(score, 1.0);
    }
}